import 'dotenv/config';
import { TextLoader } from "langchain/document_loaders/fs/text";
import { MemoryVectorStore } from "langchain/vectorstores/memory";
import { Chroma } from "@langchain/community/vectorstores/chroma";

import { RecursiveCharacterTextSplitter } from "langchain/text_splitter";
import { ChatOpenAI, OpenAIEmbeddings } from "@langchain/openai";
import { MultiQueryRetriever } from 'langchain/retrievers/multi_query';
import { ContextualCompressionRetriever } from 'langchain/retrievers/contextual_compression';
import { LLMChainExtractor } from 'langchain/retrievers/document_compressors/chain_extract';
import { ScoreThresholdRetriever } from 'langchain/retrievers/score_threshold';
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { RunnableSequence } from '@langchain/core/runnables';
import { StringOutputParser } from '@langchain/core/output_parsers';

const chatModel = new ChatOpenAI({
  configuration:{
    baseURL: process.env.OPENAI_API_BASE_URL,
  },
   apiKey: process.env.OPENAI_API_KEY,
   modelName: "doubao-1-5-pro-256k-250115",
   maxRetries: 0,

});
const embeddings = new OpenAIEmbeddings({
  configuration: {
    baseURL: process.env.OPENAI_API_BASE_URL,
  },
  apiKey: process.env.OPENAI_API_KEY,
  modelName: "doubao-embedding-text-240715",
});

const loader = new TextLoader("data/qiu.txt");
const docs = await loader.load();

const splitter = new RecursiveCharacterTextSplitter({
  chunkSize: 1000,
  chunkOverlap: 100,
});

const splitDocs = await splitter.splitDocuments(docs);

// 使用MemoryVectorStore向量存储
console.log(`开始处理 ${splitDocs.length} 个文档片段...`);

// 将所有文档加载到内存向量存储
const vectorStore = await MemoryVectorStore.fromDocuments(
  splitDocs,
  embeddings
);

console.log(`成功将 ${splitDocs.length} 个文档片段存储到内存向量存储中`);

const convertDocsToString = (documents: Document[]): string => {
  // @ts-ignore
    return documents.map((document) =>  document.pageContent).join("\n")
}

const retriever = vectorStore.asRetriever(2);



const contextRetriverChain = RunnableSequence.from([
    (input) => input.question,
    retriever,
    convertDocsToString
])

const TEMPLATE = `
你是一个熟读刘慈欣的《球状闪电》的终极原著党，精通根据作品原文详细解释和回答问题，你在回答时会引用作品原文。
并且回答时仅根据原文，尽可能回答用户问题，如果原文中没有相关内容，你可以回答“原文中没有相关内容”，

以下是原文中跟用户回答相关的内容：
{context}

现在，你需要基于原文，回答以下问题：
{question}`;

const prompt = ChatPromptTemplate.fromTemplate(
    TEMPLATE
);

const ragChain = RunnableSequence.from([
    {
        context: contextRetriverChain,
        question: (input) => input.question,
    },
    prompt,
    chatModel,
    new StringOutputParser()
])

const answer = await ragChain.invoke({
    question: "什么是球状闪电"
  });  
  console.log(answer);
